package rt2

import (
	"errors"
	"fmt"
	"net"
	"sync"

	"gitee.com/zhang-li17/rt2/pcol"

	"github.com/sirupsen/logrus"
	"google.golang.org/protobuf/proto"
)

const (
	HEADER   = 0xAA
	BUF_SIZE = 1024
)

type IMultiCastServer interface {
	Start(string, int) (error, <-chan MultiCastClient)
	Stop() error
}

// MultiCastServer 是组播服务器的实现
type MultiCastServer struct {
	// 用于存储组播连接
	multiCastConn *net.UDPConn

	// 用于停止接收消息的通道
	stop chan struct{}

	// 用于同步的互斥锁
	mu sync.Mutex

	mapLock      sync.Mutex
	clientCtxMap map[string]*connCtx

	connChan chan MultiCastClient

	running bool
}

// Start 开始监听组播地址上的连接请求
func (m *MultiCastServer) Start(ip string, port int) (<-chan MultiCastClient, error) {
	m.Stop()

	// 创建一个通道用于返回新创建的UDP连接
	m.connChan = make(chan MultiCastClient)

	// 创建一个通道用于停止接收消息
	m.stop = make(chan struct{})

	// 创建组播地址
	addr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", ip, port))
	if err != nil {
		return nil, err
	}

	// 创建组播连接
	multiCastConn, err := net.ListenMulticastUDP("udp", nil, addr)
	if err != nil {
		return nil, err
	}

	logrus.Infoln("服务器监听组播消息成功: ", addr)

	// 将组播连接存储到结构体中
	m.mu.Lock()
	m.running = true
	m.multiCastConn = multiCastConn
	m.clientCtxMap = make(map[string]*connCtx)
	m.mu.Unlock()

	// 启动一个goroutine来接收组播消息
	go m.receiveMessages(m.connChan)

	return m.connChan, nil
}

type connCtx struct {
	multicastSender *net.UDPAddr
	// 连接阶段
	connPhase int
	sync.Mutex
	name string
	port uint32
	dec  *Decoder
}

// receiveMessages 是接收组播消息的协程
func (m *MultiCastServer) receiveMessages(connChan chan<- MultiCastClient) {
	defer func() {
		// 当协程退出时，关闭组播连接
		m.mu.Lock()
		if m.multiCastConn != nil {
			m.multiCastConn.Close()
		}
		m.mu.Unlock()
	}()

	// 创建一个缓冲区用于接收消息
	buffer := make([]byte, 1024)

	for {
		// 检查是否收到停止信号
		select {
		case <-m.stop:
			return
		default:
			// 接收组播消息
			n, remoteAddr, err := m.multiCastConn.ReadFromUDP(buffer)
			if err != nil {
				fmt.Printf("Error reading from multicast: %v\n", err)
				return
			}
			m.process_read_multicast_data(remoteAddr, buffer[:n], connChan)
		}
	}
}

func (m *MultiCastServer) process_read_multicast_data(remoteAddr *net.UDPAddr, data []byte, connChan chan<- MultiCastClient) {
	// 解析消息中的客户端IP和端口
	clientStr := remoteAddr.String()

	m.mapLock.Lock()
	ctx, ok := m.clientCtxMap[clientStr]
	if !ok {
		ctx = &connCtx{
			multicastSender: remoteAddr,
			dec:             NewDecoder(HEADER, BUF_SIZE),
		}
		m.clientCtxMap[clientStr] = ctx

		m.register(ctx, connChan)
	}
	m.mapLock.Unlock()
	ctx.dec.RecvAndDecode(data)
}

func (m *MultiCastServer) register(ctx *connCtx, connChan chan<- MultiCastClient) {
	ctx.dec.RegisterDecodeCallback(FrameType(pcol.MessageType_CONNECT_REQUEST), func(f *Frame) (data any, err error) {
		req := &pcol.ConnectRequest{}
		e := proto.Unmarshal(f.Payload, req)
		if e != nil {
			logrus.Errorf("ConnectRequest from %s unmarshal error: %v\n", ctx.multicastSender.String(), e)
			return nil, e
		}

		if ctx.name == "" {
			ctx.name = req.Name
		} else if ctx.name != req.Name {
			logrus.Errorf("接收到名称变化：%s->%s", ctx.name, req.Name)
			return
		}

		// 等待连接中
		if req.Waiting {
			if ctx.connPhase != 0 {
				logrus.Errorf("phase != 0 但收到了连接请求消息，忽略")
				return
			}
			ctx.connPhase = 1

			if ctx.port == 0 {
				ctx.port = uint32(req.UnicastPort)
			} else if ctx.port != uint32(req.UnicastPort) {
				logrus.Errorf("接收到port变化: %d->%d", ctx.port, req.UnicastPort)
				return
			}

			// 连接一个UDP单播
			addr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", ctx.multicastSender.IP.String(), req.UnicastPort))
			if err != nil {
				logrus.Errorln("resolve udp addr failed: ", err)
				return nil, err
			} else {
				conn, err := net.DialUDP("udp", nil, addr)
				if err != nil {
					logrus.Errorln("dial udp addr failed: ", err)
					return nil, err
				} else {
					// 单播连接成功
					wg := sync.WaitGroup{}
					wg.Add(1)
					dec := NewDecoder(HEADER, BUF_SIZE)
					dec.RegisterDecodeCallback(FrameType(pcol.MessageType_CONNECT_CONFIRM), func(f *Frame) (any, error) {
						c := &pcol.ConnectConfirm{}
						err := proto.Unmarshal(f.Payload, c)
						if err != nil {
							logrus.Errorln("ConnectConfirm unmarshal failed: ", err)
							return nil, err
						}

						if c.LocalName == req.GetName() && c.RemoteName == "Server" {
							wg.Done()
						}
						return c, err

					})
					go func() {
						for {
							s := &pcol.ConnectConfirm{
								LocalName:  "Server",
								RemoteName: req.GetName(),
							}
							data, _ := proto.Marshal(s)
							data = dec.PackAndEncode(FrameType(pcol.MessageType_CONNECT_CONFIRM), data)
							n, err := conn.Write(data)
							if n != len(data) {
								logrus.Errorln("sending ConnectConfirm failed: ", err)
							} else {
								break
							}
						}
					}()
					buf := make([]byte, 1024)
					// conn.SetReadDeadline(time.Now().Add(time.Second * 3))
					n, raddr, err := conn.ReadFromUDP(buf)
					if err != nil {
						logrus.Errorln("等待连接确认失败:", err)
						return nil, err
					} else if raddr.String() != addr.String() {
						logrus.Errorln("等待连接确认，但接收到其它来源消息：", raddr.String())
						return nil, errors.New("等待连接确认，但接收到其它来源消息：" + raddr.String())
					}
					dec.RecvAndDecode(buf[:n])
					wg.Wait()
					logrus.Info("connect one ok:", raddr.String())
					connChan <- MultiCastClient{
						Name: req.Name,
						Conn: conn,
						Dec:  dec,
					}
				}
			}
		} else {
			logrus.Infof("客户端成功连接 %s(%s) \n", req.Name, ctx.multicastSender.String())
			m.mu.Lock()
			delete(m.clientCtxMap, ctx.multicastSender.String())
			m.mu.Unlock()
		}
		return nil, nil
	})
}

// Stop 停止组播服务器
func (m *MultiCastServer) Stop() error {
	if !m.running {
		return nil
	}

	m.mu.Lock()
	defer m.mu.Unlock()

	if m.multiCastConn == nil {
		return fmt.Errorf("server not started")
	}

	// 关闭组播连接
	if err := m.multiCastConn.Close(); err != nil {
		return err
	}

	// 发送停止信号
	close(m.stop)
	close(m.connChan)

	// 清理clientDecMap
	for clientStr := range m.clientCtxMap {
		delete(m.clientCtxMap, clientStr)
	}

	logrus.Infoln("服务器已关闭")

	m.running = false
	return nil
}

func (m *MultiCastServer) IsRunning() bool {
	return m.running
}
